"""
实现词向量和数字向量的转换
"""


class WordSequence():
    def __init__(self):
        self.UNK_TAG = "UNK"  # 不存在的词
        self.PAD_TAG = "PAD"  # 填充词

        UNK = 0
        PAD = 1

        self.dict = {  # 词典
            self.UNK_TAG: UNK,
            self.PAD_TAG: PAD
        }
        self.reverse_dict = {}  # 反转词典

    def fit(self, sentence, min_rate=None, max_rate=None, max_feature=None):
        """
        创建{"词": num}和{num: "词"}两个字典
        :param sentence: 词向量
        :param min_rate: 最小词频
        :param max_rate: 最大词频
        :param max_feature: 字典最大长度
        :return:
        """
        # 1. 将句子中的词频统计出来，放入词频字典
        word_rate_dict = {}
        for word in sentence:
            word_rate_dict[word] = word_rate_dict.get(word, 0) + 1

        # 2. 根据词频要求，修改词频字典
        if min_rate:
            word_rate_dict = {key: value for key, value in word_rate_dict.items() if value >= min_rate}
        if max_rate:
            word_rate_dict = {key: value for key, value in word_rate_dict.items() if value <= max_rate}

        # 3. 根据长度要求修改词频字典
        if max_feature:
            word_rate_dict = sorted(list(word_rate_dict.items()), key=lambda x: x[1], reverse=True)[:max_feature]

        # 4. 将词加入字典
        for word in word_rate_dict.keys():
            self.dict[word] = len(self.dict)

        # 5. 创建反转字典
        self.reverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def word_to_num_transform(self, sentence, min_length=None, max_length=None):
        """
        将词向量转换为数字向量
        :param sentence: 传入的词向量
        :param min_length: 向量最小的长度
        :param max_length: 向量的最大长度
        :return:
        """
        # 1. 遍历词向量，在字典中查找每个词对应的数字，创建句子对应的数字向量
        num_vector = []
        for word in sentence:
            num_vector.append(self.dict.get(word, self.dict[self.UNK_TAG]))

        # 2. 根据长度要求修改数字向量的长度
        if (min_length) and (len(num_vector) < min_length):
            num_vector += [self.dict[self.PAD_TAG]] * (min_length - len(num_vector))
        if (max_length) and (len(num_vector) > max_length):
            num_vector = num_vector[:max_length]

        # 3. 返回数字向量
        return num_vector

    def num_to_word_transform(self, num_vector, min_length=None, max_length=None):
        """
        将数字向量转换为词向量
        :param num_vector: 传入的数字向量
        :param min_length: 向量的最小长度
        :param max_length: 向量的最大长度
        :return:
        """
        # 1. 便利数字向量，在反转字典中查找每个数字对应的词，将词加入词向量中
        word_vector = []
        for num in num_vector:
            word_vector.append(self.reverse_dict.get(num, self.reverse_dict[0]))

        # 2. 根据长度要求修改词向量
        if (min_length) and (len(word_vector) < min_length):
            word_vector += [self.reverse_dict[1]] * (min_length - len(word_vector))
        if (max_length) and (len(word_vector) > max_length):
            word_vector = word_vector[:max_length]

        # 3. 返回词向量
        return word_vector


if __name__ == '__main__':
    wc = WordSequence()
    wc.fit(["我", "是", "谁", "呀"])
    ret1 = wc.word_to_num_transform(["是", "我", "啊"], min_length=10)
    print(ret1)

    ret2 = wc.num_to_word_transform([1, 3, 4, 2, 7], min_length=10)
    print(ret2)

